package core

import (
	"crypto/sha256"
	"github.com/hashicorp/go-hclog"
	"sync"
)

// MVBA implements Multi-valued Byzantine Agreement protocol
type MVBA struct {
	node   *Node
	logger hclog.Logger
	
	// Message storage
	valMsgs   map[int]map[int][]byte // [SN][Sender]Data
	echoMsgs  map[int]map[int][]byte // [SN][Sender]Hash
	readyMsgs map[int]map[int][]byte // [SN][Sender]Hash
	
	// Output tracking
	outputs   map[int][]byte // [SN]Data
	delivered map[int]bool   // [SN]bool
	
	// Threshold values
	t int // threshold = f+1
	n int // total nodes
	
	mutex sync.RWMutex
}

// NewMVBA creates a new MVBA instance
func NewMVBA(node *Node) *MVBA {
	return &MVBA{
		node:      node,
		logger:    hclog.New(&hclog.LoggerOptions{Name: "dumbo-mvba", Level: hclog.Level(node.Config.LogLevel)}),
		valMsgs:   make(map[int]map[int][]byte),
		echoMsgs:  make(map[int]map[int][]byte),
		readyMsgs: make(map[int]map[int][]byte),
		outputs:   make(map[int][]byte),
		delivered: make(map[int]bool),
		t:         node.F + 1,
		n:         node.N,
	}
}

// Start initiates MVBA for a given sequence number with input data
func (m *MVBA) Start(sn int, data []byte) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	
	if m.delivered[sn] {
		return // Already delivered
	}
	
	// 开始MVBA延迟计时
	m.node.delayMonitor.StartMVBA(sn)
	
	// Hash the data
	dataHash := sha256.Sum256(data)
	
	// Send Val message to all nodes
	valMsg := MVBAValMsg{
		SN:     sn,
		Sender: m.node.Id,
		Data:   data,
		Hash:   dataHash[:],
	}
	
	m.logger.Info("MVBA starting", "sn", sn, "data_len", len(data))
	
	// Broadcast val message
	m.node.PlainBroadcast(MVBAValMsgTag, valMsg, nil)
}

// HandleValMsg processes received val messages
func (m *MVBA) HandleValMsg(msg *MVBAValMsg) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	
	sn := msg.SN
	sender := msg.Sender
	data := msg.Data
	
	// Initialize maps if needed
	if m.valMsgs[sn] == nil {
		m.valMsgs[sn] = make(map[int][]byte)
	}
	
	// Store val message
	m.valMsgs[sn][sender] = data
	
	m.logger.Debug("MVBA received val", "sn", sn, "sender", sender, "val_count", len(m.valMsgs[sn]))
	
	// Check if we have t val messages
	if len(m.valMsgs[sn]) >= m.t {
		// Find the most common data
		dataCount := make(map[string]int)
		for _, d := range m.valMsgs[sn] {
			dataCount[string(d)]++
		}
		
		// Find data with at least t occurrences
		var commonData []byte
		for data, count := range dataCount {
			if count >= m.t {
				commonData = []byte(data)
				break
			}
		}
		
		if commonData != nil {
			// Send Echo message
			echoHash := sha256.Sum256(commonData)
			echoMsg := MVBAEchoMsg{
				SN:     sn,
				Sender: m.node.Id,
				Hash:   echoHash[:],
			}
			
			m.logger.Debug("MVBA sending echo", "sn", sn)
			m.node.PlainBroadcast(MVBAEchoMsgTag, echoMsg, nil)
		}
	}
}

// HandleEchoMsg processes received echo messages
func (m *MVBA) HandleEchoMsg(msg *MVBAEchoMsg) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	
	sn := msg.SN
	sender := msg.Sender
	hash := msg.Hash
	
	// Initialize maps if needed
	if m.echoMsgs[sn] == nil {
		m.echoMsgs[sn] = make(map[int][]byte)
	}
	
	// Store echo message
	m.echoMsgs[sn][sender] = hash
	
	m.logger.Debug("MVBA received echo", "sn", sn, "sender", sender, "echo_count", len(m.echoMsgs[sn]))
	
	// Check if we have t echo messages with same hash
	if len(m.echoMsgs[sn]) >= m.t {
		// Find the most common hash
		hashCount := make(map[string]int)
		for _, h := range m.echoMsgs[sn] {
			hashCount[string(h)]++
		}
		
		// Find hash with at least t occurrences
		var commonHash []byte
		for h, count := range hashCount {
			if count >= m.t {
				commonHash = []byte(h)
				break
			}
		}
		
		if commonHash != nil {
			// Send Ready message
			readyMsg := MVBAReadyMsg{
				SN:     sn,
				Sender: m.node.Id,
				Hash:   commonHash,
			}
			
			m.logger.Debug("MVBA sending ready", "sn", sn)
			m.node.PlainBroadcast(MVBAReadyMsgTag, readyMsg, nil)
		}
	}
}

// HandleReadyMsg processes received ready messages
func (m *MVBA) HandleReadyMsg(msg *MVBAReadyMsg) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	
	sn := msg.SN
	sender := msg.Sender
	hash := msg.Hash
	
	// Initialize maps if needed
	if m.readyMsgs[sn] == nil {
		m.readyMsgs[sn] = make(map[int][]byte)
	}
	
	// Store ready message
	m.readyMsgs[sn][sender] = hash
	
	m.logger.Debug("MVBA received ready", "sn", sn, "sender", sender, "ready_count", len(m.readyMsgs[sn]))
	
	// Check if we have t ready messages with same hash
	if len(m.readyMsgs[sn]) >= m.t {
		// Find the most common hash
		hashCount := make(map[string]int)
		for _, h := range m.readyMsgs[sn] {
			hashCount[string(h)]++
		}
		
		// Find hash with at least t occurrences
		var commonHash []byte
		for h, count := range hashCount {
			if count >= m.t {
				commonHash = []byte(h)
				break
			}
		}
		
		if commonHash != nil && !m.delivered[sn] {
			// Try to find the corresponding data
			if m.valMsgs[sn] != nil {
				for _, data := range m.valMsgs[sn] {
					dataHash := sha256.Sum256(data)
					if string(dataHash[:]) == string(commonHash) {
						// Deliver the data
						m.deliver(sn, data)
						break
					}
				}
			}
		}
	}
	
	// Check if we have 2t ready messages
	if len(m.readyMsgs[sn]) >= 2*m.t {
		// Find the most common hash
		hashCount := make(map[string]int)
		for _, h := range m.readyMsgs[sn] {
			hashCount[string(h)]++
		}
		
		// Find hash with at least 2t occurrences
		var commonHash []byte
		for h, count := range hashCount {
			if count >= 2*m.t {
				commonHash = []byte(h)
				break
			}
		}
		
		if commonHash != nil && !m.delivered[sn] {
			// Try to find the corresponding data
			if m.valMsgs[sn] != nil {
				for _, data := range m.valMsgs[sn] {
					dataHash := sha256.Sum256(data)
					if string(dataHash[:]) == string(commonHash) {
						// Deliver the data
						m.deliver(sn, data)
						break
					}
				}
			}
		}
	}
}

// deliver outputs the data for a given sequence number
func (m *MVBA) deliver(sn int, data []byte) {
	if m.delivered[sn] {
		return
	}
	
	m.delivered[sn] = true
	m.outputs[sn] = data
	
	// 结束MVBA延迟计时并记录
	delay := m.node.delayMonitor.EndMVBA(sn)
	m.logger.Info("MVBA delivered", "sn", sn, "data_len", len(data), "delay_ms", delay.Milliseconds())
	
	// Send output message
	outputMsg := MVBAOutputMsg{
		SN:     sn,
		Sender: m.node.Id,
		Data:   data,
	}
	
	m.node.PlainBroadcast(MVBAOutputMsgTag, outputMsg, nil)
}

// GetOutput returns the output for a given sequence number
func (m *MVBA) GetOutput(sn int) ([]byte, bool) {
	m.mutex.RLock()
	defer m.mutex.RUnlock()
	
	data, exists := m.outputs[sn]
	return data, exists
}

// IsDelivered checks if a sequence number has been delivered
func (m *MVBA) IsDelivered(sn int) bool {
	m.mutex.RLock()
	defer m.mutex.RUnlock()
	
	return m.delivered[sn]
}
